package cn.junki.sparkbigmodeldemo.listener;

import cn.hutool.json.JSONUtil;
import cn.junki.sparkbigmodeldemo.entity.request.SparkBigModelChatRequest;
import cn.junki.sparkbigmodeldemo.entity.response.SparkBigModelChatResponse;
import cn.junki.sparkbigmodeldemo.entity.vo.SparkBigModelChatVO;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.function.Consumer;

/**
 * 星火大模型WS接口监听器
 *
 * @author Junki
 * @date 2023-09-21 13:14:47
 */
@Slf4j
@NoArgsConstructor
@AllArgsConstructor
@Builder
public class SparkBigModelWebSocketListener extends WebSocketListener {

    /**
     * 参数封装
     */
    private SparkBigModelChatRequest sparkBigModelChatRequest;

    /**
     * 响应回调
     */
    private Consumer<SparkBigModelChatVO> onMessageCallBack;

    /**
     * 连接失败回调
     */
    private Consumer<Response> onFailureCallBack;

    @Override
    public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
        // 请求体
        String requestBody = JSONUtil.toJsonStr(sparkBigModelChatRequest);

        // 连接成功后请求大模型对话接口
        webSocket.send(requestBody);
        log.info("[星火大模型]请求：{}", requestBody);
    }

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
        log.info("[星火大模型]响应：{}", text);

        if (null == this.onMessageCallBack) {
            return;
        }

        // 解析响应结果
        SparkBigModelChatResponse response = JSONUtil.toBean(text, SparkBigModelChatResponse.class);

        // 错误码
        Integer code = response.getHeader().getCode();
        if (code != 0) {
            this.onMessageCallBack.accept(SparkBigModelChatVO.builder().code(code).build());
            return;
        }

        String content = response.getPayload().getChoices().getText().get(0).getContent();
        Integer status = response.getPayload().getChoices().getStatus();

        // 收到响应时回调
        this.onMessageCallBack.accept(SparkBigModelChatVO.builder().code(code).status(status).content(content).build());

        // 接收完毕后，关闭资源
        if (status == 2) {
            // https://datatracker.ietf.org/doc/html/rfc6455#section-7.4
            webSocket.close(1000, "");
        }

    }

    @Override
    public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
        log.error("[星火大模型]连接失败", t);
        if (null != this.onFailureCallBack) {
            this.onFailureCallBack.accept(response);
        }
    }
}
